﻿#include <iostream>
#include <comdef.h>
#include <Wbemidl.h>
#include <comutil.h>
#include <string>
#include <sstream>
#pragma comment(lib, "comsuppw.lib")
#pragma comment(lib, "wbemuuid.lib")
using namespace std;
BOOL wmi_run();
// BOOL wmi_getDriveLetters();
BOOL wmi_close();
string wmi_getDriveLetters(string);
IWbemLocator *pLoc = NULL;
IWbemServices *pSvc = NULL;
std::string replace(std::string str, const std::string &from, const std::string &to)
{
    std::stringstream ss;
    size_t start_pos = 0;
    while ((start_pos = str.find(from, start_pos)) != std::string::npos)
    {
        ss << str.substr(0, start_pos) << to;
        start_pos += from.length();
        str = str.substr(start_pos);
    }
    ss << str;
    return ss.str();
}
int wmi_run()
{
    HRESULT hres;

    hres = CoInitializeEx(0, COINIT_MULTITHREADED);
    if (FAILED(hres))
    {
        cout << "Failed to initialize COM library. Error code = 0x"
             << hex << hres << endl;
        return 1; // Program has failed.
    }
    hres = CoInitializeSecurity(
        NULL,
        -1,                          // COM authentication
        NULL,                        // Authentication services
        NULL,                        // Reserved
        RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication
        RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation
        NULL,                        // Authentication info
        EOAC_NONE,                   // Additional capabilities
        NULL                         // Reserved
    );

    if (FAILED(hres))
    {
        cout << "Failed to initialize security. Error code = 0x"
             << hex << hres << endl;
        CoUninitialize();
        return 1; // Program has failed.
    }

    hres = CoCreateInstance(
        CLSID_WbemLocator,
        0,
        CLSCTX_INPROC_SERVER,
        IID_IWbemLocator, (LPVOID *)&pLoc);

    if (FAILED(hres))
    {
        cout << "Failed to create IWbemLocator object."
             << " Err code = 0x"
             << hex << hres << endl;
        CoUninitialize();
        return 1; // Program has failed.
    }

    hres = pLoc->ConnectServer(
        _bstr_t(L"ROOT\\CIMV2"), // Object path of WMI namespace
        NULL,                    // User name. NULL = current user
        NULL,                    // User password. NULL = current
        0,                       // Locale. NULL indicates current
        NULL,                    // Security flags.
        0,                       // Authority (e.g. Kerberos)
        0,                       // Context object
        &pSvc                    // pointer to IWbemServices proxy
    );

    if (FAILED(hres))
    {
        cout << "Could not connect. Error code = 0x"
             << hex << hres << endl;
        pLoc->Release();
        CoUninitialize();
        return 1; // Program has failed.
    }
    hres = CoSetProxyBlanket(
        pSvc,                        // Indicates the proxy to set
        RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
        RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
        NULL,                        // Server principal name
        RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx
        RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
        NULL,                        // client identity
        EOAC_NONE                    // proxy capabilities
    );

    if (FAILED(hres))
    {
        cout << "Could not set proxy blanket. Error code = 0x"
             << hex << hres << endl;
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        return 1; // Program has failed.
    }
    return 0;
}

BOOL wmi_close()
{

    pSvc->Release();
    pLoc->Release();
    CoUninitialize();

    return 0;
}

int wmi_getDriveLetters_str(string DRS)
{
    DRS = DRS.substr(0, 2);
    HRESULT hres;
    IEnumWbemClassObject *pEnumerator = NULL;
    hres = pSvc->ExecQuery(
        bstr_t("WQL"),
        bstr_t("SELECT * FROM Win32_DiskDrive"),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        NULL,
        &pEnumerator);

    if (FAILED(hres))
    {
        cout << "Query for processes failed. "
             << "Error code = 0x"
             << hex << hres << endl;
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        return FALSE; // Program has failed.
    }
    else
    {

        IWbemClassObject *pclsObj;
        ULONG uReturn = 0;
        while (pEnumerator)
        {
            pEnumerator->Next(WBEM_INFINITE, 1,
                                     &pclsObj, &uReturn);
            if (0 == uReturn)
                break;
            VARIANT vtProp;
            pclsObj->Get(_bstr_t(L"DeviceID"), 0, &vtProp, 0, 0);
            wstring tmp = vtProp.bstrVal;
            tmp = tmp.substr(4);
            pclsObj->Get(L"Caption", 0, &vtProp, 0, 0);
            wstring wstrQuery = L"Associators of {Win32_DiskDrive.DeviceID='\\\\.\\";
            wstrQuery += tmp;
            wstrQuery += L"'} where AssocClass=Win32_DiskDriveToDiskPartition";
            IEnumWbemClassObject *pEnumerator1 = NULL;
            hres = pSvc->ExecQuery(
                bstr_t("WQL"),
                bstr_t(wstrQuery.c_str()),
                WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
                NULL,
                &pEnumerator1);

            if (FAILED(hres))
            {
                cout << "Query for processes failed. "
                     << "Error code = 0x"
                     << hex << hres << endl;
                pSvc->Release();
                pLoc->Release();
                CoUninitialize();
                return FALSE; // Program has failed.
            }
            else
            {

                IWbemClassObject *pclsObj1;
                ULONG uReturn1 = 0;
                while (pEnumerator1)
                {
                    pEnumerator1->Next(WBEM_INFINITE, 1,
                                              &pclsObj1, &uReturn1);
                    if (0 == uReturn1)
                        break;
                    VARIANT vtProp1;
                    pclsObj1->Get(_bstr_t(L"DeviceID"), 0, &vtProp1, 0, 0);
                    // printf("DeviceID:%ls\n", vtProp1.bstrVal);
                    wstring wstrQuery = L"Associators of {Win32_DiskPartition.DeviceID='";
                    wstrQuery += vtProp1.bstrVal;
                    wstrQuery += L"'} where AssocClass=Win32_LogicalDiskToPartition";
                    IEnumWbemClassObject *pEnumerator2 = NULL;
                    hres = pSvc->ExecQuery(
                        bstr_t("WQL"),
                        bstr_t(wstrQuery.c_str()),
                        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
                        NULL,
                        &pEnumerator2);
                    if (FAILED(hres))
                    {
                        cout << "Query for processes failed. "
                             << "Error code = 0x"
                             << hex << hres << endl;
                        pSvc->Release();
                        pLoc->Release();
                        CoUninitialize();
                        return FALSE; // Program has failed.
                    }
                    else
                    {
                        IWbemClassObject *pclsObj2;
                        ULONG uReturn2 = 0;
                        while (pEnumerator2)
                        {
                            pEnumerator2->Next(WBEM_INFINITE, 1,
                                                      &pclsObj2, &uReturn2);
                            if (0 == uReturn2)
                                break;
                            VARIANT vtProp2;
                             pclsObj2->Get(_bstr_t(L"DeviceID"), 0, &vtProp2, 0, 0);
                            // print result

                            if (DRS == _com_util::ConvertBSTRToString(vtProp2.bstrVal))
                            {
                                string strls = replace(_com_util::ConvertBSTRToString(vtProp1.bstrVal), "Disk #", "");
                                strls = replace(strls, " Partition #", "");
                                strls = replace(strls, ",", ":");
                                printf("硬盘ID:%s\n", strls.c_str());
                                //  Disk #1  Partition #2
                            }
                            VariantClear(&vtProp2);
                        }
                        pclsObj1->Release();
                    }
                    VariantClear(&vtProp1);
                    pEnumerator2->Release();
                }
                pclsObj->Release();
            }
            VariantClear(&vtProp);
            pEnumerator1->Release();
        }
    }
    pEnumerator->Release();
    return 1;
}